from osgeo import gdal
import numpy as np
import matplotlib.pyplot as plt
import os


def read_raster(path_img,band):
    os.getcwd()
    ## 读入栅格数据，包括影像信息和地理信息
    dset = gdal.Open(path_img)
    geo_trans = dset.GetGeoTransform()   ### 获取地理转换参数：（左上角坐标x, 像元宽度，行旋转参数，左上角坐标y,列旋转参数，像元高度）
    print('geo_transform:', geo_trans)
    proj = dset.GetProjection()    ### 获取影像投影
    print('projection:', proj)
    x_size, y_size, num_band = dset.RasterXSize, dset.RasterYSize, dset.RasterCount  ## 获取影像尺寸、波段数信息
    print('xsize: %.0f, ysize: %.0f, num_band: %.0f' % (x_size, y_size, num_band))
    band = dset.GetRasterBand(band)   ### 获取影像某一个波段
    band_array = band.ReadAsArray()     ### 将影像读入为np.array()格式
    print('shape of array:', band_array.shape)
    return band_array

def read_trans(path_img):
        dset = gdal.Open(path_img)
        geo_trans = dset.GetGeoTransform()





def write_raster(path,band_array,geo_trans,proj):
    driver = gdal.GetDriverByName("GTiff")  ## 创建.tif文件驱动
    outdset = driver.Create(path, xsize=band_array.shape[1], \
                                    ysize=band_array.shape[0], bands=1, eType=gdal.GDT_Int16)  ### 创建空的.tif数据
    outdset.SetGeoTransform(geo_trans)    ### 设置地理转换参数
    outdset.SetProjection(proj)           ### 设置投影
    outband = outdset.GetRasterBand(1)    ### 获取波段1
    outband.WriteArray(band_array)        ### 将np.array()数组写入波段1
    outband.SetNoDataValue(np.nan)        ### 设置特定值为无数据，具有压缩功能
    outdset = None    ## 关闭.tif文件驱动
    print("done")